"""
负荷预测工作流演示脚本
展示如何使用可插拔的数据获取函数进行Transformer负荷预测
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import datetime
import sys
import os
from typing import Callable
import warnings
warnings.filterwarnings('ignore')

# 导入工作流
from load_prediction_workflow import HVACLoadPredictionWorkflow, create_workflow, quick_train_and_predict

# 导入原始数据获取函数
try:
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from System_of_C.read_and_process_data.read_data_for_Loadpredict import calculate_load
    REAL_DATA_AVAILABLE = True
except ImportError:
    print("⚠️ 无法导入真实数据获取函数，将使用模拟数据")
    REAL_DATA_AVAILABLE = False


def create_mock_data_function() -> Callable:
    """
    创建模拟数据获取函数
    模拟真实环境中的数据格式和特征
    """
    def mock_calculate_load(start_time: str, end_time: str) -> pd.DataFrame:
        """
        模拟数据获取函数，返回格式与真实函数一致
        
        Args:
            start_time: 开始时间字符串
            end_time: 结束时间字符串
            
        Returns:
            DataFrame: 包含timestamp, swt, swh, real_load的数据
        """
        print(f"📊 模拟数据获取: {start_time} 至 {end_time}")
        
        start_dt = pd.to_datetime(start_time)
        end_dt = pd.to_datetime(end_time)
        
        # 生成小时级时间序列
        timestamps = pd.date_range(start_dt, end_dt, freq='H')[:-1]  # 不包含结束时间
        n_points = len(timestamps)
        
        if n_points == 0:
            print("⚠️ 时间范围无效，返回空数据")
            return pd.DataFrame(columns=['timestamp', 'swt', 'swh', 'real_load'])
        
        # 模拟负荷数据（考虑日周期和周周期）
        hours = np.array([ts.hour for ts in timestamps])
        days = np.array([ts.dayofyear for ts in timestamps])
        weekdays = np.array([ts.weekday() for ts in timestamps])
        
        # 基础负荷模式（日周期：白天高，夜晚低）
        base_load = 100 + 80 * np.sin(2 * np.pi * hours / 24 - np.pi/2)
        
        # 周周期（工作日vs周末）
        weekend_factor = np.where(weekdays >= 5, 0.7, 1.0)  # 周末负荷降低
        
        # 季节性变化
        season_factor = 1 + 0.3 * np.sin(2 * np.pi * days / 365 - np.pi/2)
        
        # 添加随机噪声和异常值
        noise = np.random.normal(0, 8, n_points)
        
        # 偶尔的设备异常（负荷突然变化）
        anomaly_indices = np.random.choice(n_points, size=max(1, n_points//100), replace=False)
        anomaly_factors = np.ones(n_points)
        anomaly_factors[anomaly_indices] = np.random.uniform(0.3, 1.8, len(anomaly_indices))
        
        real_load = base_load * weekend_factor * season_factor * anomaly_factors + noise
        real_load = np.maximum(real_load, 10)  # 确保负荷为正值
        
        # 模拟室外温度（考虑日变化和季节变化）
        daily_temp_variation = 8 * np.sin(2 * np.pi * hours / 24 - np.pi/2)
        seasonal_temp = 20 + 15 * np.sin(2 * np.pi * days / 365 - np.pi/2)
        swt = seasonal_temp + daily_temp_variation + np.random.normal(0, 2, n_points)
        
        # 模拟含湿量（与温度相关，夏季高冬季低）
        base_humidity = 0.008 + 0.006 * np.sin(2 * np.pi * days / 365 - np.pi/2)
        humidity_noise = np.random.normal(0, 0.001, n_points)
        swh = base_humidity + humidity_noise
        swh = np.clip(swh, 0.002, 0.025)  # 限制含湿量范围
        
        # 创建DataFrame
        data = pd.DataFrame({
            'timestamp': timestamps,
            'swt': swt,
            'swh': swh,
            'real_load': real_load
        })
        
        print(f"✅ 模拟数据生成完成: {len(data)} 个时间点")
        print(f"   负荷范围: {real_load.min():.1f} - {real_load.max():.1f}")
        print(f"   温度范围: {swt.min():.1f} - {swt.max():.1f}°C")
        print(f"   含湿量范围: {swh.min():.4f} - {swh.max():.4f}")
        
        return data
    
    return mock_calculate_load


def demo_1_basic_workflow():
    """演示1: 基本工作流使用"""
    print("\n" + "="*60)
    print("🚀 演示1: 基本工作流使用")
    print("="*60)
    
    # 创建模拟数据函数
    mock_data_func = create_mock_data_function()
    
    # 创建工作流
    workflow = create_workflow(mock_data_func)
    
    # 测试数据获取和处理
    start_time = "2024-01-01 00:00:00"
    end_time = "2024-01-20 00:00:00"  # 2天数据用于测试
    
    try:
        processed_data, raw_data = workflow.fetch_and_process_data(start_time, end_time)
        
        print(f"✅ 数据处理成功")
        print(f"📊 原始数据: {len(raw_data)} 条记录")
        print(f"🔧 处理后特征数量: {len(processed_data)}")
        for key, value in processed_data.items():
            if isinstance(value, np.ndarray):
                print(f"   {key}: {value.shape}")
        
    except Exception as e:
        print(f"❌ 数据处理失败: {e}")


def demo_2_custom_data_function():
    """演示2: 使用不同的数据获取函数"""
    print("\n" + "="*60)
    print("🔄 演示2: 切换数据获取函数")
    print("="*60)
    
    # 创建第一个数据函数（高负荷场景）
    def high_load_scenario(start_time: str, end_time: str) -> pd.DataFrame:
        mock_func = create_mock_data_function()
        data = mock_func(start_time, end_time)
        # 将负荷提高50%模拟高负荷场景
        data['real_load'] *= 1.5
        print("📈 高负荷场景数据")
        return data
    
    # 创建第二个数据函数（低负荷场景）
    def low_load_scenario(start_time: str, end_time: str) -> pd.DataFrame:
        mock_func = create_mock_data_function()
        data = mock_func(start_time, end_time)
        # 将负荷降低30%模拟低负荷场景
        data['real_load'] *= 0.7
        print("📉 低负荷场景数据")
        return data
    
    # 测试不同数据函数
    workflow = create_workflow()
    start_time = "2024-01-01 00:00:00"
    end_time = "2024-01-20 00:00:00"
    
    scenarios = [
        ("高负荷场景", high_load_scenario),
        ("低负荷场景", low_load_scenario)
    ]
    
    for name, func in scenarios:
        print(f"\n🔄 切换到: {name}")
        workflow.set_data_fetch_function(func)
        
        try:
            _, raw_data = workflow.fetch_and_process_data(start_time, end_time)
            avg_load = raw_data['real_load'].mean()
            print(f"   平均负荷: {avg_load:.2f}")
        except Exception as e:
            print(f"   ❌ 失败: {e}")


def demo_3_training_and_prediction():
    """演示3: 模型训练和预测"""
    print("\n" + "="*60)
    print("🎯 演示3: 模型训练和预测")
    print("="*60)
    
    # 使用更长的时间范围进行训练
    start_time = "2024-01-01 00:00:00"
    end_time = "2024-03-15 00:00:00"  # 2周数据
    
    mock_data_func = create_mock_data_function()
    
    try:
        print("🚀 开始快速训练和预测...")
        result = quick_train_and_predict(
            start_time, end_time,
            mock_data_func,
            epochs=10  # 减少epoch数量以加快演示
        )
        
        print(f"\n✅ 训练和预测完成!")
        print(f"📊 训练结果:")
        print(f"   验证损失: {result['training_result']['best_val_loss']:.6f}")
        print(f"   验证MAE: {result['training_result']['final_val_mae']:.6f}")
        
        print(f"🔮 预测结果:")
        stats = result['prediction_stats']
        print(f"   预测时长: {stats['remaining_hours']} 小时")
        print(f"   平均负荷: {stats['mean_load']:.2f}")
        print(f"   峰值负荷: {stats['max_load']:.2f}")
        print(f"   最低负荷: {stats['min_load']:.2f}")
        
        return result
        
    except Exception as e:
        print(f"❌ 训练或预测失败: {e}")
        return None


def demo_4_real_data_integration():
    """演示4: 真实数据集成（如果可用）"""
    print("\n" + "="*60)
    print("🌍 演示4: 真实数据集成")
    print("="*60)
    
    if not REAL_DATA_AVAILABLE:
        print("⚠️ 真实数据获取函数不可用，跳过此演示")
        return
    
    # 使用真实数据获取函数
    workflow = create_workflow(calculate_load)
    
    # 测试真实数据获取
    start_time = "2025-08-15 00:00:00"
    end_time = "2025-09-25 01:00:00"
    
    try:
        processed_data, raw_data = workflow.fetch_and_process_data(start_time, end_time)
        
        print(f"✅ 真实数据获取成功")
        print(f"📊 数据统计:")
        print(f"   记录数量: {len(raw_data)}")
        print(f"   负荷范围: {raw_data['real_load'].min():.2f} - {raw_data['real_load'].max():.2f}")
        print(f"   温度范围: {raw_data['swt'].min():.2f} - {raw_data['swt'].max():.2f}°C")
        print(f"   含湿量范围: {raw_data['swh'].min():.4f} - {raw_data['swh'].max():.4f}")
        
        # 检查数据质量
        missing_data = raw_data.isnull().sum()
        if missing_data.sum() > 0:
            print(f"⚠️ 发现缺失数据:")
            for col, count in missing_data.items():
                if count > 0:
                    print(f"   {col}: {count} 个缺失值")
        
    except Exception as e:
        print(f"❌ 真实数据获取失败: {e}")
        print("💡 可能的原因: 数据库连接问题、时间范围无数据、权限问题等")


def demo_5_visualization():
    """演示5: 预测结果可视化"""
    print("\n" + "="*60)
    print("📈 演示5: 预测结果可视化")
    print("="*60)
    
    # 先进行一次快速训练
    mock_data_func = create_mock_data_function()
    
    try:
        # 生成测试数据
        start_time = "2024-01-01 00:00:00"
        end_time = "2024-01-08 00:00:00"
        
        print("🚀 快速训练模型用于可视化...")
        result = quick_train_and_predict(
            start_time, end_time,
            mock_data_func,
            epochs=5
        )
        
        workflow = result['workflow']
        predictions = result['predictions']
        
        # 获取测试数据用于对比
        test_start = "2024-01-07 00:00:00"
        test_end = "2024-01-08 00:00:00"
        _, test_data = workflow.fetch_and_process_data(test_start, test_end)
        
        # 绘制预测结果
        plt.figure(figsize=(15, 8))
        
        # 子图1: 历史数据和预测
        plt.subplot(2, 2, 1)
        # 绘制历史数据（最后24小时）
        recent_data = test_data.tail(24)
        plt.plot(range(len(recent_data)), recent_data['real_load'], 
                'b-', label='历史负荷', linewidth=2)
        
        # 绘制预测数据
        pred_start_idx = len(recent_data)
        pred_indices = range(pred_start_idx, pred_start_idx + len(predictions))
        plt.plot(pred_indices, predictions, 'r--', label='预测负荷', linewidth=2)
        
        plt.title('负荷预测对比')
        plt.xlabel('时间点')
        plt.ylabel('负荷')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 子图2: 预测统计
        plt.subplot(2, 2, 2)
        stats = result['prediction_stats']
        categories = ['平均值', '最大值', '最小值']
        values = [stats['mean_load'], stats['max_load'], stats['min_load']]
        bars = plt.bar(categories, values, color=['blue', 'red', 'green'], alpha=0.7)
        plt.title('预测负荷统计')
        plt.ylabel('负荷')
        
        # 在柱状图上添加数值标签
        for bar, value in zip(bars, values):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(values)*0.01,
                    f'{value:.1f}', ha='center', va='bottom')
        
        # 子图3: 环境数据
        plt.subplot(2, 2, 3)
        plt.plot(recent_data['swt'], label='室外温度', color='orange')
        plt.title('环境温度变化')
        plt.xlabel('时间点')
        plt.ylabel('温度 (°C)')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 子图4: 含湿量
        plt.subplot(2, 2, 4)
        plt.plot(recent_data['swh'], label='含湿量', color='purple')
        plt.title('含湿量变化')
        plt.xlabel('时间点')
        plt.ylabel('含湿量')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        # 保存图片
        save_path = 'load_prediction_demo.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"📊 可视化结果已保存为: {save_path}")
        plt.show()
        
    except Exception as e:
        print(f"❌ 可视化演示失败: {e}")


def main():
    """主演示函数"""
    print("🏢 基于Transformer的空调负荷预测工作流演示")
    print("支持可插拔数据获取函数，便于不同数据源集成")
    
    try:
        # 运行所有演示
        demo_1_basic_workflow()
        demo_2_custom_data_function()
        
        # 询问是否运行训练演示（耗时较长）
        user_input = input("\n是否运行训练和预测演示？(可能需要几分钟) [y/N]: ")
        if user_input.lower() in ['y', 'yes']:
            result = demo_3_training_and_prediction()
            
            if result is not None:
                # 询问是否运行可视化
                viz_input = input("\n是否显示预测结果可视化？ [y/N]: ")
                if viz_input.lower() in ['y', 'yes']:
                    demo_5_visualization()
        
        demo_4_real_data_integration()
        
        print("\n" + "="*60)
        print("✅ 所有演示完成！")
        print("="*60)
        
        print("\n🎯 工作流主要特点:")
        print("✓ 支持可插拔的数据获取函数")
        print("✓ 自动数据预处理和特征工程")
        print("✓ 基于Transformer的高精度预测")
        print("✓ 实时预测和批量预测能力")
        print("✓ 完整的训练到部署流程")
        print("✓ 灵活的配置和扩展性")
        
        print("\n📖 使用建议:")
        print("1. 根据数据源实现相应的数据获取函数")
        print("2. 确保数据函数返回标准格式: timestamp, swt, swh, real_load")
        print("3. 根据数据量和计算资源调整模型配置")
        print("4. 定期重新训练模型以适应新的数据模式")
        
    except KeyboardInterrupt:
        print("\n\n⚠️ 演示被用户中断")
    except Exception as e:
        print(f"\n❌ 演示过程中发生错误: {e}")
        print("💡 请检查依赖库安装和环境配置")


if __name__ == "__main__":
    main()
